/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.fn.harness;



import com.google.auto.service.AutoService;

import com.bff.gaia.unified.fn.harness.control.BundleSplitListener;

import com.bff.gaia.unified.fn.harness.data.UnifiedFnDataClient;

import com.bff.gaia.unified.fn.harness.data.PCollectionConsumerRegistry;

import com.bff.gaia.unified.fn.harness.data.PTransformFunctionRegistry;

import com.bff.gaia.unified.fn.harness.state.UnifiedFnStateClient;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi.RemoteGrpcPort;

import com.bff.gaia.unified.model.pipeline.v1.Endpoints;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi.Components;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi.PCollection;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi.PTransform;

import com.bff.gaia.unified.runners.core.construction.CoderTranslation;

import com.bff.gaia.unified.runners.core.construction.RehydratedComponents;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.fn.data.CloseableFnDataReceiver;

import com.bff.gaia.unified.sdk.fn.data.FnDataReceiver;

import com.bff.gaia.unified.sdk.fn.data.LogicalEndpoint;

import com.bff.gaia.unified.sdk.fn.data.RemoteGrpcPortWrite;

import com.bff.gaia.unified.sdk.options.PipelineOptions;

import com.bff.gaia.unified.sdk.util.WindowedValue;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableMap;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import java.io.IOException;

import java.util.Map;

import java.util.function.Supplier;



import static com.bff.gaia.unified.vendor.guava.com.google.common.collect.Iterables.getOnlyElement;



/**

 * Registers as a consumer with the Unified Fn Data Api. Consumes elements and encodes them for

 * transmission.

 *

 * <p>Can be re-used serially across {@link UnifiedFnApi.ProcessBundleRequest}s. For each request, call

 * {@link #registerForOutput()} to start and call {@link #close()} to finish.

 */

public class UnifiedFnDataWriteRunner<InputT> {



  private static final Logger LOG = LoggerFactory.getLogger(UnifiedFnDataWriteRunner.class);



  /** A registrar which provides a factory to handle writing to the Fn Api Data Plane. */

  @AutoService(PTransformRunnerFactory.Registrar.class)

  public static class Registrar implements PTransformRunnerFactory.Registrar {



    @Override

    public Map<String, PTransformRunnerFactory> getPTransformRunnerFactories() {

      return ImmutableMap.of(RemoteGrpcPortWrite.URN, new Factory());

    }

  }



  /** A factory for {@link UnifiedFnDataWriteRunner}s. */

  static class Factory<InputT> implements PTransformRunnerFactory<UnifiedFnDataWriteRunner<InputT>> {



    @Override

    public UnifiedFnDataWriteRunner<InputT> createRunnerForPTransform(

        PipelineOptions pipelineOptions,

        UnifiedFnDataClient unifiedFnDataClient,

        UnifiedFnStateClient unifiedFnStateClient,

        String pTransformId,

        PTransform pTransform,

        Supplier<String> processBundleInstructionId,

        Map<String, PCollection> pCollections,

        Map<String, RunnerApi.Coder> coders,

        Map<String, RunnerApi.WindowingStrategy> windowingStrategies,

        PCollectionConsumerRegistry pCollectionConsumerRegistry,

        PTransformFunctionRegistry startFunctionRegistry,

        PTransformFunctionRegistry finishFunctionRegistry,

        BundleSplitListener splitListener)

        throws IOException {

      UnifiedFnApi.Target target =

          UnifiedFnApi.Target.newBuilder()

              .setPrimitiveTransformReference(pTransformId)

              .setName(getOnlyElement(pTransform.getInputsMap().keySet()))

              .build();

      RunnerApi.Coder coderSpec;

      if (RemoteGrpcPortWrite.fromPTransform(pTransform).getPort().getCoderId().isEmpty()) {

        LOG.error(

            "Missing required coder_id on grpc_port for %s; using deprecated fallback.",

            pTransformId);

        coderSpec =

            coders.get(

                pCollections.get(getOnlyElement(pTransform.getInputsMap().values())).getCoderId());

      } else {

        coderSpec = null;

      }

      UnifiedFnDataWriteRunner<InputT> runner =

          new UnifiedFnDataWriteRunner<>(

              pTransform, processBundleInstructionId, target, coderSpec, coders, unifiedFnDataClient);

      startFunctionRegistry.register(pTransformId, runner::registerForOutput);

      pCollectionConsumerRegistry.register(

          getOnlyElement(pTransform.getInputsMap().values()),

          pTransformId,

          (FnDataReceiver) (FnDataReceiver<WindowedValue<InputT>>) runner::consume);



      finishFunctionRegistry.register(pTransformId, runner::close);

      return runner;

    }

  }



  private final Endpoints.ApiServiceDescriptor apiServiceDescriptor;

  private final UnifiedFnApi.Target outputTarget;

  private final Coder<WindowedValue<InputT>> coder;

  private final UnifiedFnDataClient unifiedFnDataClientFactory;

  private final Supplier<String> processBundleInstructionIdSupplier;



  private CloseableFnDataReceiver<WindowedValue<InputT>> consumer;



  UnifiedFnDataWriteRunner(

      RunnerApi.PTransform remoteWriteNode,

      Supplier<String> processBundleInstructionIdSupplier,

      UnifiedFnApi.Target outputTarget,

      RunnerApi.Coder coderSpec,

      Map<String, RunnerApi.Coder> coders,

      UnifiedFnDataClient unifiedFnDataClientFactory)

      throws IOException {

    RemoteGrpcPort port = RemoteGrpcPortWrite.fromPTransform(remoteWriteNode).getPort();

    this.apiServiceDescriptor = port.getApiServiceDescriptor();

    this.unifiedFnDataClientFactory = unifiedFnDataClientFactory;

    this.processBundleInstructionIdSupplier = processBundleInstructionIdSupplier;

    this.outputTarget = outputTarget;



    RehydratedComponents components =

        RehydratedComponents.forComponents(Components.newBuilder().putAllCoders(coders).build());

    @SuppressWarnings("unchecked")

	Coder<WindowedValue<InputT>> coder;

    if (!port.getCoderId().isEmpty()) {

      coder =

          (Coder<WindowedValue<InputT>>)

              CoderTranslation.fromProto(coders.get(port.getCoderId()), components);

    } else {

      // TODO: remove this path once it is no longer used

      coder = (Coder<WindowedValue<InputT>>) CoderTranslation.fromProto(coderSpec, components);

    }

    this.coder = coder;

  }



  public void registerForOutput() {

    consumer =

        unifiedFnDataClientFactory.send(

            apiServiceDescriptor,

            LogicalEndpoint.of(processBundleInstructionIdSupplier.get(), outputTarget),

            coder);

  }



  public void close() throws Exception {

    consumer.close();

  }



  public void consume(WindowedValue<InputT> value) throws Exception {

    consumer.accept(value);

  }

}